#!/usr/bin/env python

# Script to go through and build a pickled dictionary archive of
# all the GFS ensemble members high, low, qpf and p12 forecasts
# for verification purposes and to determine weights
import sys
sys.path.append('../data_parsers')
from itertools import groupby
from datetime import datetime, timedelta
import re, os
import numpy as np
from numpy import nan
from scipy.stats.stats import nanmean
from encode_truth import encode_truth
import pickle


# Get the station ID we want from the command line
from optparse import OptionParser
parser = OptionParser()
parser.add_option('-s', '--site', dest='siteid', action='store', type='string', default='KSEA', help='4-letter id of desired site')
(opts,args) = parser.parse_args()
siteid = opts.siteid.upper()

cldcov = {'CL' : 0,
          'FW' : 1,
          'SC' : 2,
          'PC' : 3,
          'BK' : 4,
          'OV' : 5}
          


def main():
    truthdir = encode_truth(siteid)
    err_archive = build_archive(siteid, truthdir)
    pickle.dump(err_archive, open('%s_gfse_archive.pickle' % siteid,'w'))
    #err_archive = pickle.load(open('%s_gfse_archive.pickle'% siteid,'r'))
    stdict = compute_stats(err_archive)
    biasdict = bias_plot(err_archive)    
    write_stats(stdict,biasdict, siteid)


def write_stats(stdict,biasdict,siteid):
    # Combine the dicts into one and then go
    print stdict[0].keys()
    for mem in stdict.keys():
        stdict[mem].update(biasdict[mem])
    outdict = stdict
    pickle.dump(outdict, open('%s_gfse_stats.pickle' % siteid, 'w'))



def compute_stats(err_archive):
    # Loop through each ensemble member and 
    # compute stats about the error
    statdict = {}
    for mem in err_archive.keys():
        statdict[mem] = {}
        statdict[mem]['hierr'] = []
        statdict[mem]['loerr'] = []
        statdict[mem]['closest_hi'] = []
        statdict[mem]['closest_lo'] = []

    for daten in err_archive[0].keys():
        hilist = []
        lolist = []
        for mem in range(0,21):
            hilist.append(abs(err_archive[mem][daten]['hierr']))
            lolist.append(abs(err_archive[mem][daten]['loerr']))
            statdict[mem]['hierr'].append(err_archive[mem][daten]['hierr'])
            statdict[mem]['loerr'].append(err_archive[mem][daten]['loerr'])
        for mem in range(0,21):
            if mem == hilist.index(min(hilist)):
                statdict[mem]['closest_hi'].append(1)
            else:
                statdict[mem]['closest_hi'].append(0)
            if mem == lolist.index(min(lolist)):
                statdict[mem]['closest_lo'].append(1)
            else:    
                statdict[mem]['closest_lo'].append(0)

    # Now can compute statistics
    outdict = {}
    mem_weights = np.zeros(len(statdict.keys()))
    for mem in statdict.keys():
        outdict[mem] = {}
        for dset in statdict[mem].keys():
            if not dset.startswith('closest'):
                outdict[mem][dset] = {}
                dlist = statdict[mem][dset]
                bias = np.mean(dlist)
                spread = np.std(dlist)
                abs_list = [abs(d) for d in dlist]
                mae = np.mean(abs_list)
                outdict[mem][dset]['bias'] = bias
                outdict[mem][dset]['spread'] = spread
                outdict[mem][dset]['mae'] = mae

                print "Mem %d %s --> Bias: %f Std: %f MAE: %f" % (mem,dset,bias,spread,mae)
            else:
                dlist = statdict[mem][dset]
                dsum = np.sum(dlist)
                dlen = len(dlist)
                outdict[mem][dset] = float(dsum)/dlen
                print "Member", mem, dset, dsum, dlen, float(dsum)/dlen
    hi_wt_list = [outdict[mem]['closest_hi'] for mem in range(len(statdict.keys()))]
    lo_wt_list = [outdict[mem]['closest_lo'] for mem in range(len(statdict.keys()))]
    plot_weights(hi_wt_list,'High Temperature')
    plot_weights(lo_wt_list,'Low Temperature')
    return outdict


def plot_weights(weightlist, var):
    # Do a simple bar chart of the weights
    import matplotlib.pyplot as plt
    plt.figure()
    cenvals = [r - 0.5 for r in range(len(weightlist))]
    plt.bar(cenvals,weightlist,width = 0.8,bottom=0)
    plt.title('Relative weights of each member for %s' % var)
    plt.show()
    


def bias_plot(err_archive):
    import matplotlib.pyplot as plt
    outdict = {}
    plt.figure()
    for mem in err_archive.keys():
        outdict[mem] = {}
        print "Member", mem
        model_his = []
        actual_his = []
        #model_los = []
        #actual_los = []
        for daten in err_archive[mem].keys():
            model_his.append(err_archive[mem][daten]['high'])
            actual_his.append(err_archive[mem][daten]['high'] + err_archive[mem][daten]['hierr'])
            #model_los.append(err_archive[mem][daten]['low'])
            #actual_los.append(err_archive[mem][daten]['low'] + err_archive[mem][daten]['loerr'])
        xvals = range(int(min(model_his)),int(max(model_his)))
        w = np.linalg.lstsq(np.array([model_his, np.ones(len(model_his))]).T,np.array(actual_his))[0]
        line = [w[0] * x + w[1] for x in xvals]
        print np.shape(line), np.shape(xvals)
        outdict[mem]['hibias_x0'] = w[1]
        outdict[mem]['hibias_x1'] = w[0]
        plt.subplot(3,7,1+mem)
        plt.scatter(model_his,actual_his)
        plt.hold(True)
        plt.plot(xvals,line,'r-')
        plt.title('Member %d' % mem)
        plt.xlabel('x0=%1.3f x1=%1.3f' % (w[1],w[0]))
        plt.hold(False)
    plt.show()

    plt.figure()
    for mem in err_archive.keys():
        print "Member", mem
        #model_his = []
        #actual_his = []
        model_los = []
        actual_los = []
        for daten in err_archive[mem].keys():
            #model_his.append(err_archive[mem][daten]['high'])
            #actual_his.append(err_archive[mem][daten]['high'] + err_archive[mem][daten]['hierr'])
            model_los.append(err_archive[mem][daten]['low'])
            actual_los.append(err_archive[mem][daten]['low'] + err_archive[mem][daten]['loerr'])
        xvals = range(int(min(model_los)),int(max(model_los)))
        w = np.linalg.lstsq(np.array([model_los, np.ones(len(model_los))]).T,np.array(actual_los))[0]
        line = [w[0] * x + w[1] for x in xvals]
        print np.shape(line), np.shape(xvals)
        outdict[mem]['lobias_x0'] = w[1]
        outdict[mem]['lobias_x1'] = w[0]
        plt.subplot(3,7,1+mem)
        plt.scatter(model_los,actual_los)
        plt.hold(True)
        plt.plot(xvals,line,'r-')
        plt.title('Member %d' % mem)
        plt.xlabel('x0=%1.3f x1=%1.3f' % (w[1],w[0]))
        plt.hold(False)
    plt.show()
    return outdict

def build_archive(siteid,truthdir):
    members = {}
    for m in range(0,21):
        members[m] = {}


    for file in os.listdir('./data'): #['data/ensp01mex.200704']:
        print file
        pmem = int(re.search('(\d{2})mex',file).groups()[0])

        for empty, record in groupby(open('./data/'+file,'r'),key=str.isspace):
            if not empty:
                bklines = list(record)
                if bklines[0].strip().startswith(siteid):
                    #print "Found!"
                    #print bklines
                    # Get the forecast date
                    analysis_date = re.search('(\d{1,2}/\d{1,2}/\d{4})',bklines[0]).groups()[0]
                    analysis_dt = datetime.strptime(analysis_date,'%m/%d/%Y')
                    fcst_date = analysis_dt + timedelta(hours=24)
                    #print "DATE FOUND:", analysis_dt
                    dates = bklines[1].split('|')
                    forecast_dts = []
                    for date in dates:
                        datesp = date.split()
                        fcst_time = analysis_dt + timedelta(hours=int(datesp[-1])-24)
                        forecast_dts.append(fcst_time)
                    #print forecast_dts


                    # Now we have a list of dates.  Use this to key the different variables
                    # Start with temperatures
                    xn = bklines[3].split('|')
                    xn[0] = xn[0][4:]
                    num = xn[forecast_dts.index(fcst_date)]
                    if len(num.split()) == 2:
                        low, high = num.split()
                        if int(low) == 999:
                            low = nan 
                        else:
                            low = int(low)
                        if int(high) == 999:
                            high = int(high)
                        else:
                            high = int(high)
                    elif len(num.split()) == 1:
                        if int(num) == 999:
                            high = nan 
                        else:
                            high = int(num)
                        low = nan
                    else:
                        high = nan
                        low = nan
                    
                    forecast_index = forecast_dts.index(fcst_date)
                    # Use a function to get the mean of each of these values
                    dpt = get_var(bklines,5,forecast_index)
                    cld = get_var(bklines,6,forecast_index)
                    p24 = get_var(bklines,8,forecast_index)
                    q24 = get_var(bklines,10,forecast_index)
                   

                    actual_high = truthdir[fcst_date]['high']
                    actual_low = truthdir[fcst_date]['low']
                    high_error = actual_high - high
                    low_error = actual_low - low

                    #print "Forecast for %s %d %d | %d %d" % (fcst_date.strftime('%m/%d/%Y'),high,low,actual_high,actual_low)
                    members[pmem][fcst_date] = {}
                    members[pmem][fcst_date]['hierr'] = high_error
                    members[pmem][fcst_date]['loerr'] = low_error
                    members[pmem][fcst_date]['high'] = high
                    members[pmem][fcst_date]['low'] = low
                    members[pmem][fcst_date]['dpt'] = dpt
                    members[pmem][fcst_date]['cld'] = cld
                    members[pmem][fcst_date]['p24'] = p24 
                    members[pmem][fcst_date]['q24'] = q24
    return members
                    

def get_var(blockln,linenum,fcst_index):
    # Now split the variable and return the mean
    var = blockln[linenum].split('|')
    var[0] = var[0][4:]
    num = var[fcst_index]

    # Special treatment for cloud cover
    if linenum == 6:
       cldcodes = [str(cldcov[r]) for r in num.split()]
       #print "CLDCODES", cldcodes
       num = ' '.join(cldcodes)

    if len(num.split()) == 2:
        var1,var2 = num.split()
        if int(var1) == 999:
            var1 = nan 
        else:
            var1 = int(var1)
        if int(var2) == 999:
            var2 = int(var2)
        else:
            var2 = int(var2)
    elif len(num.split()) == 1:
        if int(num) == 999:
            var2 = nan 
        else:
            var2 = int(num)
            var1 = nan
    else:
        var2 = nan
        var1 = nan

    meanvar = nanmean([var1,var2])
    return meanvar



if __name__ == '__main__':
    main()
